import streamlit as st
import os
import cv2
import numpy as np
from multiprocessing import Pool
from functools import partial

def process_image(image_path, output_dir):
    """改进的图像处理函数"""
    try:
        # 使用imdecode替代imread来处理中文路径
        img = cv2.imdecode(np.fromfile(image_path, dtype=np.uint8), cv2.IMREAD_GRAYSCALE)
        
        # 检查图像是否成功读取
        if img is None:
            st.error(f"无法读取图像: {image_path}")
            return None
            
        # 检查图像尺寸
        if img.size == 0:
            st.error(f"图像为空: {image_path}")
            return None
        
        # 使用自适应阈值代替固定阈值
        img_bin = cv2.adaptiveThreshold(img, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
                                      cv2.THRESH_BINARY, 11, 2)
        
        # 使用边缘检测改进裁剪区域识别
        edges = cv2.Canny(img_bin, 100, 200)
        
        # 使用轮廓检测找到最大连通区域
        contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        if not contours:
            st.error(f"未找到有效轮廓: {image_path}")
            return None
            
        largest_contour = max(contours, key=cv2.contourArea)
        x, y, w, h = cv2.boundingRect(largest_contour)
        
        # 裁剪图像
        cropped = img[y:y+h, x:x+w]
        
        # 保存处理后的图像
        output_path = os.path.join(output_dir, os.path.basename(image_path))
        # 使用imencode替代imwrite来处理中文路径
        cv2.imencode('.bmp', cropped)[1].tofile(output_path)
        
        return output_path
        
    except Exception as e:
        st.error(f"处理图像时出错 {image_path}: {str(e)}")
        return None

def preview_processing(image_path):
    """预览处理效果"""
    try:
        # 处理路径中的特殊字符
        image_path = os.path.normpath(image_path).replace('\\', '/')
        
        # 检查文件是否存在
        if not os.path.exists(image_path):
            st.error(f"文件不存在: {image_path}")
            return None, None
            
        # 使用完整的二进制模式读取文件
        with open(image_path, 'rb') as f:
            file_bytes = np.asarray(bytearray(f.read()), dtype=np.uint8)
            img = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR)
            
        if img is None:
            st.error(f"无法读取图像: {image_path}")
            return None, None
            
        # 其余代码保持不变
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        img_bin = cv2.adaptiveThreshold(gray, 255,
                                      cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
                                      cv2.THRESH_BINARY, 11, 2)
        edges = cv2.Canny(img_bin, 100, 200)
        return img, edges
        
    except Exception as e:
        st.error(f"处理图像时出错: {str(e)}")
        return None, None

def main():
    st.title("PIV图像防抖处理系统")
    
    # 侧边栏设置
    st.sidebar.header("设置")
    input_dir = st.sidebar.text_input("输入目录", 
        "I:/数据合计/zhongwenyi/棒束小量程下游较早期视频-订正无误已填表(分类完毕)/离散泡状流/jf=1.7, jg=0.019_C001H001S0001pao")
    
    # 规范化路径
    input_dir = os.path.normpath(input_dir).replace('\\', '/')
    
    # 添加详细的路径检查
    if not os.path.exists(input_dir):
        st.error(f"输入目录不存在: {input_dir}")
        st.write("请检查路径是否正确，特别是逗号和空格的位置")
        return
        
    # 显示当前目录内容以便调试
    try:
        files = os.listdir(input_dir)
        st.write("目录中的文件：", files[:5])  # 只显示前5个文件
    except Exception as e:
        st.error(f"读取目录出错: {str(e)}")
        return
    
    output_dir = st.sidebar.text_input("输出目录", "输出文件夹路径")
    output_dir = output_dir.replace('\\', '/')
    num_cores = st.sidebar.slider("使用CPU核心数", 1, os.cpu_count(), os.cpu_count())
    
    # 文件选择
    if st.sidebar.button("选择文件"):
        files = [f for f in os.listdir(input_dir) if f.endswith('.bmp')]
        if not files:
            st.error("所选目录中没有找到.bmp文件")
            return
            
        selected_file = st.sidebar.selectbox("选择预览文件", files)
        
        if selected_file:
            # 显示预览
            full_path = os.path.join(input_dir, selected_file)
            img, edges = preview_processing(full_path)
            if img is not None and edges is not None:
                col1, col2 = st.columns(2)
                with col1:
                    st.image(img, caption="原始图像", use_column_width=True)
                with col2:
                    st.image(edges, caption="边缘检测结果", use_column_width=True)
    
    # 运行按钮
    if st.sidebar.button("开始处理"):
        # 获取所有bmp文件
        files = [os.path.join(input_dir, f) for f in os.listdir(input_dir) 
                if f.endswith('.bmp')]
        
        if not files:
            st.error("未找到.bmp文件")
            return
            
        # 创建输出目录
        os.makedirs(output_dir, exist_ok=True)
        
        # 使用进度条
        progress_bar = st.progress(0)
        
        # 使用多核并行处理
        with Pool(num_cores) as pool:
            results = []
            for i, result in enumerate(pool.imap(partial(process_image, output_dir=output_dir), files)):
                results.append(result)
                # 更新进度条
                progress_bar.progress((i + 1) / len(files))
        
        # 过滤掉None结果
        successful_results = [r for r in results if r is not None]
        st.success(f"处理完成！成功处理{len(successful_results)}/{len(files)}张图片。")

if __name__ == "__main__":
    main()